import time
import pandas as pd
import sqlite3
import numpy as np
import torch

def get_inverse_block(mis, mat, n):
    # return Sig_12 Sig22^{-1}, Sig_11 - Sig12 Sig_22^{-1} SIg21
    start_time = time.time()
    print(mis)
    sig11 = mat[mis,mis]
    sig12 = torch.cat([mat[mis,:mis], mat[mis,(mis+1):]])
    sig22 = torch.cat([torch.cat([mat[:mis, :mis], mat[:mis,(mis+1):]], axis=1),
    torch.cat([mat[(mis+1):, :mis], mat[(mis+1):,(mis+1):]], axis=1)])
    sig22_inverse = torch.cholesky_inverse(torch.linalg.cholesky(sig22))
    print(time.time() - start_time)
    temp = sig12 @ sig22_inverse
    return (temp), sig11 - (temp@ sig12)


def LoadData(DataName, OutcomeCol):
    DB = sqlite3.connect('Database/analysisdata.db')
    dt = pd.read_sql_query("SELECT * from " + DataName, DB)
    dt = dt.set_index('usid')
    DB.close()
    # always change 0-1 encodig for 'econ_c', 'admin_c', 'health_c', 'soccap_c'
    # Therefore, outcome = 1 always means good result.
    for i in range(6):
        if np.corrcoef(dt.my4, dt.iloc[:,-6+i])[0,1] < -0.05:
            dt.iloc[:,-6+i] = 1 - dt.iloc[:,-6+i]
    # sort by safety score decreasing
    dt = dt.sort_values(by='my4', ascending=False)
    X = torch.Tensor(np.array(dt.iloc[:, :19]))
    A = torch.Tensor(np.array(dt['my4'])).to(torch.int)
    Y = torch.Tensor(np.array(dt[OutcomeCol]).reshape((-1, 1)))
    return X,A,Y

